import sys
sys.path.append('./')

import argparse
import yaml
from pathlib import Path
from easydict import EasyDict as edict

import torch
import torch.onnx
from utils import *

parser = argparse.ArgumentParser(description='PyTorch Image Classification Training')
parser.add_argument('config', default='configs/res18.yaml', type=str, nargs='?', help='config file path')
parser.add_argument('--resume', type=str, help='ckpt file path')

if __name__ == '__main__':
    cmd_args = parser.parse_args()
    args = yaml.load(open(cmd_args.config), yaml.Loader)
    args = edict(args)

    out_file = Path(cmd_args.config).stem+'.onnx'

    model = build_model(args.model)
    resume_from_ckpt(model, None, cmd_args.resume, device=torch.device('cpu'))

    x = torch.randn([1, 3, 224, 224])
    torch.onnx.export(
        model, x, out_file,
        training=torch.onnx.TrainingMode.TRAINING,
        opset_version=12,
        do_constant_folding=False
    )